#!/usr/bin/env python3
# -*- coding:utf-8 -*-
"""
生成 SegmentMinCoo AclNN 轻量化验收所需的单用例数据
"""

import argparse
from pathlib import Path

import numpy as np

try:
    import torch
    from torch_scatter import segment_coo

    USE_SCATTER = True
except Exception:  # pragma: no cover
    torch = None  # type: ignore
    segment_coo = None  # type: ignore
    USE_SCATTER = False

SRC_SHAPE = (8, 2048)
INDEX_SHAPE = (8,)
DEFAULT_SEGMENTS = 4
DEFAULT_CASE = "single_case"
DTYPE = np.float32


def write_binary(path: Path, array: np.ndarray) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    array.tofile(path)
    print(f"[INFO] 写入 {path} | shape={array.shape}, dtype={array.dtype}")


def compute_segment_min(src: np.ndarray, index: np.ndarray, n_segments: int):
    """按照 SegmentMin 语义生成 out 与 arg_out"""
    rows, cols = src.shape
    arg_out = np.full((n_segments, cols), -1, dtype=np.int32)

    if USE_SCATTER:
        torch_src = torch.from_numpy(src)
        torch_index = torch.from_numpy(index.astype(np.int64))
        valid_mask = (torch_index >= 0) & (torch_index < n_segments)

        if torch.any(valid_mask):
            valid_rows = torch.nonzero(valid_mask, as_tuple=True)[0]
            filtered_src = torch_src.index_select(0, valid_rows)
            filtered_index = torch_index.index_select(0, valid_rows)
            scatter_out = segment_coo(
                filtered_src, filtered_index,dim_size=n_segments, reduce="min"
            )
            out = scatter_out.cpu().numpy()
        else:
            out = np.zeros((n_segments, cols), dtype=src.dtype)
    else:
        out = np.full((n_segments, cols), np.inf, dtype=src.dtype)

        # Pass1: 统计 segment 最小值
        for row in range(rows):
            seg = int(index[row])
            if seg < 0 or seg >= n_segments:
                continue
            out[seg] = np.minimum(out[seg], src[row])

    non_empty = np.zeros(n_segments, dtype=bool)

    # Pass2: 记录最小值对应的 logicalM 索引（同 kernel 行为，最后一次覆盖）
    for row in range(rows):
        seg = int(index[row])
        if seg < 0 or seg >= n_segments:
            continue
        mask = np.isclose(src[row], out[seg], rtol=1e-6, atol=1e-6)
        arg_out[seg][mask] = row
        non_empty[seg] = True

    # 空 segment 按算子约定置零
    for seg in range(n_segments):
        if not non_empty[seg]:
            out[seg].fill(0.0)

    return out.astype(DTYPE), arg_out


def main():
    parser = argparse.ArgumentParser(description="生成 SegmentMinCoo 单用例数据")
    parser.add_argument("--case", type=str, default=DEFAULT_CASE, help="用例名称，默认 single_case")
    parser.add_argument("--seed", type=int, default=2025, help="随机种子，默认 2025")
    parser.add_argument("--nsegments", type=int, default=DEFAULT_SEGMENTS, help="segment 数量，默认 4")
    parser.add_argument("--base-dir", type=str, default="..", help="输入输出根目录，默认上级目录")
    args = parser.parse_args()

    rng = np.random.default_rng(args.seed)
    src = rng.uniform(-1.0, 1.0, size=SRC_SHAPE).astype(DTYPE)
    # 固定 pattern，便于复现
    base_pattern = np.arange(INDEX_SHAPE[0], dtype=np.int32) % args.nsegments
    index = np.copy(base_pattern)

    out, arg_out = compute_segment_min(src, index, args.nsegments)

    base_dir = Path(__file__).resolve().parent
    input_dir = base_dir / "input"
    output_dir = base_dir / "output"

    case = args.case

    write_binary(input_dir / f"src_{case}.bin", src)
    write_binary(input_dir / f"index_{case}.bin", index.astype(np.int32))
    write_binary(output_dir / f"golden_{case}.bin", out)
    write_binary(output_dir / f"golden_arg_out_{case}.bin", arg_out.astype(np.int32))
    print(f"src: {src}")
    print(f"index: {index}")
    print(f"out: {out}")
    print(f"arg_out: {arg_out}")
    print(f"[INFO] 用例 {case} 数据生成完毕")


if __name__ == "__main__":
    main()

